%-------------------------------------------------------------------------------
% $Id: vorticity_stream.m,v 1.8 2011/07/19 16:17:30 paul Exp $
% $Date: 2011/07/19 16:17:30 $
% $Author: paul $
%-------------------------------------------------------------------------------
% Vorticity Stream Function Equations
% Arguments:  
%             graphic (str):  'imagesc'(default), 'quiver'
%             k (float):       time step increment, default is 0.2
%             M (int):         grid axis points, default is 64
%             nu (float):      fluid viscosity, default is 0.0002
%             mode (str):      pattern of initial vortices at t=0
%                              'dual' (default), 'quad', 'five', 'bar',
%                              'opposite_balanced', 'opposite_unbalanced',
%
% Coupled Equations:   zeta{t} = -(u*zeta){x} -(v*zeta){y}   convection
%                                + nu*(zeta{xx} + zeta{yy})  diffusion
%                      zeta = psi{xx} + psi{yy}              vorticity
%                      u = psi{y}, v = -psi{x}               velocities
% Domain:              -1 =< x,y =< +1 
%                      0 =< t =< 90 
%                      viscosity nu = 0.0002
% Initial Conditions:  two vortices in the x,y plane at t=0 (default mode)
%                      described by the following equation:
%                      zeta(x,y,t) = exp(-20*((x-0.5)^2+y^2)) 
%                                    + exp(-20*((x+0.5)^2+y^2))
% Boundary Conditions: u(-1,y,t) = u(1,y,t) = u(x,-1,t) = u(x,1,t) = 0
% 
% Method:
% 0. We are given the initial conditions on zeta as well as
%    boundary conditions
% 1. Solve the Poisson Equation for psi
% 2. Calculate velocities u and v from psi
% 3. Use the Convection-Diffusion equation to advance zeta to the
%    next time step. Repeat steps 1 to 3 until all time steps done.
%
% Example Invocations:
% vorticity_stream
% vorticity_stream('quiver',0.2,64,0.0002,'dual')
% vorticity_stream('imagesc',0.2,64,0.0002,'five')
% vorticity_stream('imagesc',0.1,64,0.0002,'bar')
% vorticity_stream('quiver',0.2,64,0.0002,'opposite_unbalanced')
%-------------------------------------------------------------------------------
function vorticity_stream(graphic,k,M,nu,mode)
if nargin < 1
  graphic = 'imagesc';
  k = 0.2;  % time increment dt, unstable when > 0.5  
  M = 64;   % number of grid points on x and y axis  
  nu = 0.0002;  % viscosity  
  mode = 'dual';
elseif nargin < 2
  k = 0.2;  
  M = 64;   
  nu = 0.0002; 
  mode = 'dual';
elseif nargin < 3
  M = 64;   
  nu = 0.0002;
elseif nargin < 4
  nu = 0.0002;
  mode = 'dual';
elseif nargin < 5
  mode = 'dual';
end

%T = 90; 
T=200;
N = ceil(T/k);

% create mesh
x = linspace(-1,1,M+2);
y = x;
h = 2/(M+2);

st_crit1 = (nu * k)/(h^2 + h^2);
if st_crit1 > 0.5
  msg = sprintf('WARNING:stability criterium VIOLATED: %f <= 0.5',st_crit1);    
  disp(msg); 
end


%-------------------------------------------------------------------------------
% Setup AVI file and parameters
aviobj = avifile('vorticity_stream.avi');
aviobj.compression = 'None';   % no other option for Unix
aviobj.fps = floor(N/30);      % frames per second (scaled to 30 seconds)
aviobj.quality = 15;           % low quality is OK

% function handle selects plot function
if strcmp(graphic,'imagesc')
  do_plot = str2func('do_imagesc');
else
  do_plot = str2func('do_quiver');                           
end
fig = figure;
%-------------------------------------------------------------------------------
tic;
mask = make_mask(M+2);
zeta = zeros(M+2,M+2); 
zeta = initial_vorticity(x(:),y(:),mode);

c = sparse_pentadiagonal(-4.0,1.0,1.0,1.0,1.0,(M+2));  

for t=1:N % time update loop
  % solve Poisson equation for the stream function psi
  zetav = reshape(zeta,(M+2)^2,1);  % convert zeta matrix to vector
  psiv = -(h^2).*(c\zetav);     % stream function psi 
  psi = reshape(psiv,M+2,M+2);  % vector back to matrix

  % calculate velocities: u = psi{y} and v = -psi{x}
  u = zeros(M,M); v = zeros(M,M);
  u =  (psi(2:M+1,2:M+1) - psi(2:M+1,1:M))./h;
  v = -(psi(2:M+1,2:M+1) - psi(1:M,2:M+1))./h;

  % this stability test only applies for standard finitie
  % differences scheme and not the upwind-approximation scheme
  %  uv_stability_limit = 100;                 % arbitrary value
  %  uv_stability(u,v,h,nu,uv_stabilit_limit); % determine u,v stability limit
  
  % calculate midpoint velocities for upwind algorithm
  % u interpolated in x direction
  uR = 0.5 * (u + [u(2:M,:) ; u(M,:)]);
  uL = 0.5 * (u + [u(1,:)   ; u(1:M-1,:)]);
  % v interpolated in y direction
  vT = 0.5 * (v + [v(:,2:M)   v(:,M)]);
  vB = 0.5 * (v + [v(:,1)     v(:,1:M-1)]);

  % upwind-difference approximation for improved stability
  % stability is still not guaranteed using this scheme!
  zetaR = zeros(M+2); zetaL = zeros(M+2); zetaT = zeros(M+2); zetaB = zeros(M+2);
  for i=2:M+1
    i1 = i-1;    
    for j=2:M+1  
      j1 = j-1;
      
      if uR(i1,j1) >=0
        zetaR(i,j)=zeta(i,j);
      else
        zetaR(i,j)=zeta(i+1,j);
      end
      if uL(i1,j1) < 0
        zetaL(i,j)=zeta(i,j);      
      else
        zetaL(i,j)=zeta(i-1,j);              
      end
      if vT(i1,j1) >=0
        zetaT(i,j)=zeta(i,j);
      else
        zetaT(i,j)=zeta(i,j+1);
      end
      if vB(i1,j1) < 0
        zetaB(i,j)=zeta(i,j);      
      else
        zetaB(i,j)=zeta(i,j-1);              
      end
    end
  end

  % calculate the convection term
  convection = -(uR.*zetaR(2:M+1,2:M+1) -uL.*zetaL(2:M+1,2:M+1));
  convection = convection - (vT.*zetaT(2:M+1,2:M+1) -vB.*zetaB(2:M+1,2:M+1));
  convection = convection ./h;
  convection_pad = pad_matrix(convection); % size up to M+2
  
  % calculate diffusion term
  diffusionv = (nu/(h^2)) * (c*zetav);
  diffusion = reshape(diffusionv,M+2,M+2);
 
  zeta = ((k*(convection_pad + diffusion)) + zeta) .* mask;

  do_plot(x,y,v,u,zeta);
  title('2D Vorticity Streamfunction Solution');
  frame = getframe(fig);  
  aviobj = addframe(aviobj,frame);
end

time_elapsed = toc;                % timing stopwatch off
msg = sprintf('INFO:\tElapsed Runtime: %d seconds', floor(time_elapsed));
disp(msg);

% Clean up rendering frame and AVI file
close(fig);
aviobj = close(aviobj);

%-------------------------------------------------------------------------------
% Functions
%-------------------------------------------------------------------------------
function matrix = sparse_pentadiagonal(main,p1,m1,p2,m2,m)
% construct sparse pentadiagonal matrix 
% zero boundary condition is compensated by inserting
% zero elements on the first (positive and negative)
% off diagonals at regular intervals of (i*m) and (i*m)+1

a = zeros(m^2,m^2);
p1v = repmat(p1,m^2,1);
m1v = repmat(m1,m^2,1);
mainv = repmat(main,m^2,1);
p2v = repmat(p2,m^2,1);
m2v = repmat(m2,m^2,1);

% columns of b represent diagonal bands in constructed matrix
b = [m2v m1v mainv p1v p2v]; 
d = [-m, -1, 0, 1, m];   % diagonal band relative indices
matrix = spdiags(b,d,a);

% insert zeros in first off-diagonals for boundary condition
for i=2:m^2-1
  mod_m = mod(i,m);
  if mod_m == 0
    matrix(i,i+1)=0;
  elseif mod_m == 1
    matrix(i,i-1)=0;
  end
end

function matrix = initial_vorticity(x,y,mode)
% create initial, time=0 vorticity matrix zeta
M = size(x,1);
matrix = zeros(M,M); 
for i=1:M
  for j=1:M
    xi = x(i); yj = y(j);
    if strcmp(mode,'dual')
      matrix(i,j) = exp(-20*((xi-0.5).^2+yj.^2)) + exp(-20*((xi+0.5).^2+yj.^2));
    elseif strcmp(mode,'quad')
      matrix(i,j) = exp(-40*((xi-0.5).^2+(yj-0.5).^2)) ... 
          + exp(-40*((xi-0.5).^2+(yj+0.5).^2)) ... 
          + exp(-40*((xi+0.5).^2+(yj-0.5).^2)) ... 
          + exp(-40*((xi+0.5).^2+(yj+0.5).^2));      
    elseif strcmp(mode,'opposite_balanced')
      matrix(i,j) = exp(-20*((xi-0.5).^2+yj.^2)) - exp(-20*((xi+0.5).^2+yj.^2));    
    elseif strcmp(mode,'opposite_unbalanced')
      matrix(i,j) = 2* exp(-20*((xi-0.5).^2+yj.^2)) - exp(-20*((xi+0.5).^2+yj.^2));        
    elseif strcmp(mode,'bar')    
      if i>(M/16) & i<(15*M/16) & j>(8*M/16) & j<(9*M/16)
        matrix(i,j)=1;
      end
    elseif strcmp(mode,'negative_dual')        
      matrix(i,j) = -exp(-20*((xi-0.5).^2+yj.^2)) - exp(-20*((xi+0.5).^2+yj.^2));      
    elseif strcmp(mode,'five')
      matrix(i,j) = exp(-20*(xi.^2+yj.^2)) ... 
          + exp(-30*((xi-0.5).^2+(yj-0.5).^2)) ...           
          + exp(-30*((xi-0.5).^2+(yj+0.5).^2)) ... 
          + exp(-30*((xi+0.5).^2+(yj-0.5).^2)) ... 
          + exp(-30*((xi+0.5).^2+(yj+0.5).^2));      
    end
  end
end

function scaled_quiver(factor,v,u)
  vr = downsample(v,factor);
  vc = downsample(vr',factor);
  vd = vc';
  ur = downsample(u,factor);
  uc = downsample(ur',factor);
  ud = uc';
  M = size(vr,1);
  quiver(1:M,1:M,vd,ud); 

function do_imagesc(x,y,v,u,z)
  imagesc(rot90(z), [0 0.25]); colormap hot;

function do_quiver(x,y,v,u,z)
  scaled_quiver(4,v,u);
  
function mask = make_mask(M)
  % construct square boundary mask with the edges set to zero
  mask = zeros(M);                 
  mask(2:M-1,2:M-1) = ones(M-2);         

function mpad = pad_matrix(m)
  msize = size(m,1);
  mpad = zeros(msize+2);
  mpad(2:msize+1,2:msize+1) = m;

function uv_stability(u,v,h,nu,limit)
  % for each element of u and v check the criterium
  % max(abs(u .* dx), abs(v .* dy)) / nu <= limit, limit=2
  % This criterium is for the standard convection algorithm and may
  % be relaxed when using the upwind approximation algorithm
  criteria_matrix = max(limit, (max(abs(u),abs(v)).*(h/nu))) - limit*ones(size(u));
  violations = size(find(criteria_matrix));
  if violations > 0
    msg = sprintf('WARNING:stability criterium VIOLATED on %d (u,v) points', violations);        
    disp(msg); 
  end

    